# -*- coding: utf-8 -*-
# ftp_server.py
# 这是一个FTP服务器，使用socket方式编程，从创建socket、监听端口开始，实现FTP协议的功能
# 导入所需的模块
import os
from time import sleep

import psutil
from PyQt5 import QtWidgets, QtCore, QtGui
from PyQt5.QtCore import pyqtSignal, QObject, QThread
import sys
import socket
import threading
import platform

import db_manager
from PyQt5.QtCore import pyqtSignal, QObject

# 定义一些常量
HOST = '127.0.0.1' # FTP服务器的IP地址，可以修改为其他值
PORT = 8888 # FTP服务器的端口号，可以修改为其他值
BUFFER_SIZE = 1024 # 缓冲区大小，用于接收和发送数据
COMMANDS = ['ls', 'cd', 'get', 'put', 'restart', 'login', 'register', 'quit'] # 支持的FTP命令
BASE_DIR = os.path.dirname(os.path.abspath(__file__)) # FTP服务器的根目录，可以修改为其他值

# 定义一个FTP服务器类
class FTPServer(QObject):
    # 定义一个信号，用于传递日志消息
    log_signal = pyqtSignal(str)
    # 初始化方法
    def __init__(self, HOST, PORT):
        super().__init__()
        # 增加一个属性，用于存储断点的位置
        self.breakpoint = 0
        self.lock = threading.Lock()
        # 创建一个socket对象，用于监听客户端的连接
        self.server_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        # 设置socket的选项，允许重用地址，避免端口占用的问题
        self.server_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        # 绑定IP地址和端口号
        self.server_sock.bind((HOST, PORT))
        # 开始监听，设置最大连接数为5
        self.server_sock.listen(5)
        # 打印服务器启动的消息
        # print('FTP服务器启动，监听地址：', HOST, ':', PORT)
        print(self.log_signal)
        self.log_signal.emit('FTP服务器启动，监听地址：' + HOST + ':' + str(PORT))
        print("信号已发送")
        # 控制关闭handle线程
        self.receive_file_event = threading.Event()
        self.receive_file_event.set()
        # 去除数据头大小标识
        self.size_flag = 0
        self.user_role = ''

    # 启动服务器的方法
    def start(self):
        # 循环接受客户端的连接
        while True:
            # 接受客户端的连接，返回一个客户端的socket对象和地址
            client_sock, client_addr = self.server_sock.accept()
            # 打印客户端连接的消息
            self.log_signal.emit('客户端连接：' + str(client_addr))
            # 创建一个子线程，用于处理客户端的请求
            threading.Thread(target=self.handle_client, args=(client_sock, client_addr)).start()

    # 处理客户端的请求的方法
    def handle_client(self, client_sock, client_addr):
        # 在子线程中创建一个DBManager对象
        db = db_manager.DBManager()
        # 发送一个欢迎消息给客户端
        client_sock.send('欢迎使用FTP服务器'.encode())
        # 初始化客户端的当前目录为服务器的根目录
        current_dir = BASE_DIR
        # 循环接收客户端的命令
        while True:
            # 等待 receive_file 事件被设置
            self.receive_file_event.wait()
            # 尝试接收客户端的命令
            try:
                # 接收客户端的命令
                command = client_sock.recv(BUFFER_SIZE).decode()
                # 如果命令为空，就跳过本次循环
                if not command:
                    continue
                # 在控制台打印客户端的命令
                self.log_signal.emit('接收命令：' + str(command))
                # 如果命令不是支持的FTP命令，就发送一个错误消息给客户端
                if command.split(' ')[0] not in COMMANDS:
                    client_sock.send('错误的命令'.encode())
                    continue
                # 根据不同的命令，执行不同的操作
                if command.startswith('login'):
                    # 如果是login命令，就处理登录请求
                    username, password = command.split(' ')[1:]
                    if db.query_user(username, password):
                        self.user_role = db.get_user_role(username, password)
                        client_sock.send(f"OK 登录成功，角色: {self.user_role}".encode())
                    else:
                        client_sock.send("ERROR 用户名或密码错误".encode())
                    continue
                if command == 'ls':
                    # 如果是ls命令，就发送当前目录和文件列表给客户端
                    self.list_dir(client_sock, current_dir)
                elif command.startswith('cd'):
                    # 如果是cd命令，就切换当前目录，并发送结果给客户端
                    current_dir = self.change_dir(client_sock, command, current_dir)
                elif command.startswith('get'):
                    # 如果是get命令，就向客户端发送文件
                    self.send_file(client_sock, command, current_dir)
                elif command.startswith('put'):
                    # 如果是put命令，且是管理员就可以接收客户端发送文件
                    if self.user_role == "管理员":
                        # 清除事件，暂停 handle_client 线程
                        self.receive_file_event.clear()
                        # 如果是put命令，就创建一个子线程，把接收文件的方法作为目标函数，把服务器的响应作为参数
                        t = threading.Thread(target=self.receive_file, args=(client_sock, command, current_dir))
                        t.start()
                        # 等待接收文件线程结束
                        t.join()
                        # 设置事件，恢复 handle_client 线程
                        self.receive_file_event.set()
                        # self.set_breakpoint(client_sock, 'restart 0')
                    else:
                        client_sock.send("ERROR 无权限写入文件".encode())
                elif command.startswith('restart'):
                    # 如果是restart命令，就设置断点
                    self.set_breakpoint(client_sock, command)
                elif command.startswith('register'):
                    # 如果是register命令，就处理注册请求
                    self.add_user_to_database(client_sock, command, db)
                elif command == 'quit':
                    # 如果是quit命令，就关闭客户端的socket，退出循环
                    client_sock.close()
                    break
            # 如果发生异常，就关闭客户端的socket，退出循环
            except Exception as e:
                self.log_signal.emit('客户端断开：'+str(client_addr))
                client_sock.close()
                break

    # 发送当前目录和文件列表给客户端的方法
    def list_dir(self, client_sock, current_dir):
        osystem = platform.system()
        self.log_signal.emit("操作系统：" + osystem)
        if osystem == 'Windows':
            # 如果当前目录是\\，就列出所有磁盘
            if current_dir == '\\':
                # 获取所有磁盘的名称
                drives = os.popen('wmic logicaldisk get name').read().split()
                # 去掉列表中的第一个元素，它是一个标题
                drives.pop(0)
                # 把当前目录和磁盘列表拼接成一个字符串，用换行符分隔
                response = current_dir + '\n' + '\n'.join(drives)
            # 否则，就列出当前目录下的所有文件和文件夹
            else:
                # 获取当前目录下的所有文件和文件夹
                files = os.listdir(current_dir)
                # 创建一个空列表，用于存储加上\的目录名和文件大小
                dir_files = []
                # 循环遍历文件列表
                for file in files:
                    # 拼接当前目录和文件名，得到文件的完整路径
                    filepath = os.path.join(current_dir, file)
                    # 如果文件是一个目录，就在文件名后面加上\
                    if os.path.isdir(filepath):
                        file += '\\'
                    # 否则，就获取文件的大小，以字节为单位
                    else:
                        size = os.path.getsize(filepath)
                        file = str(size) + ' ' + file
                    # 把文件名添加到列表中
                    dir_files.append(file)
                # 把当前目录和文件列表拼接成一个字符串，用换行符分隔
                response = current_dir + '\n' + '\n'.join(dir_files)
            # 发送响应给客户端
            client_sock.send(response.encode())
        elif osystem == 'Linux':
            # 如果当前目录是/，就列出所有磁盘
            if current_dir == '/':
                return
            # 获取当前目录下的所有文件和文件夹
            files = os.listdir(current_dir)
            # 创建一个空列表，用于存储加上/的目录名和文件大小
            dir_files = []
            # 循环遍历文件列表
            for file in files:
                # 拼接当前目录和文件名，得到文件的完整路径
                filepath = os.path.join(current_dir, file)
                # 如果文件是一个目录，就在文件名后面加上/
                if os.path.isdir(filepath):
                    file += '/'
                # 否则，就获取文件的大小，以字节为单位
                else:
                    size = os.path.getsize(filepath)
                    file = str(size) + ' ' + file
                # 把文件名添加到列表中
                dir_files.append(file)
            # 把当前目录和文件列表拼接成一个字符串，用换行符分隔
            response = current_dir + '\n' + '\n'.join(dir_files)
            # 发送响应给客户端
            client_sock.send(response.encode())

    # 切换当前目录并发送结果给客户端的方法
    def change_dir(self, client_sock, command, current_dir):
        osystem = platform.system()
        self.log_signal.emit("操作系统：" + osystem)
        if osystem == 'Windows':
            # 把命令分割为两部分，第一部分是cd，第二部分是目标目录
            parts = command.split(' ', 1)
            if len(parts) >= 2:
                _, target_dir = parts
            else:
                # 如果没有目标目录，默认为根目录
                target_dir = '/'
            # 如果目标目录是..，就返回上一级目录
            if target_dir == '..':
                # 如果当前目录是磁盘的根目录，就返回一个特殊的目录，表示所有磁盘
                if current_dir.endswith(':\\'):
                    current_dir = '\\'
                # 否则，就返回上一级目录
                else:
                    current_dir = os.path.dirname(current_dir.rstrip('\\'))
            # 否则，就拼接当前目录和目标目录，得到新的目录
            else:
                # 如果目标目录不以\结尾，就加上\
                if not target_dir.endswith('\\'):
                    target_dir += '\\'
                current_dir = os.path.join(current_dir, target_dir)
            # 如果新的目录存在，就发送一个成功的响应给客户端
            if os.path.exists(current_dir):
                response = 'OK ' + current_dir
            # 否则，就发送一个失败的响应给客户端
            else:
                response = '目录不存在'
            # 发送响应给客户端
            client_sock.send(response.encode())
            # 返回新的当前目录
            return current_dir
        elif osystem == 'Linux':
            # 把命令分割为两部分，第一部分是cd，第二部分是目标目录
            parts = command.split(' ', 1)
            if len(parts) >= 2:
                _, target_dir = parts
            else:
                # 如果没有目标目录，默认为根目录
                target_dir = '/'
            # 如果目标目录是..，就返回上一级目录
            if target_dir == '..':
                # 返回上一级目录
                current_dir = os.path.dirname(current_dir.rstrip('/'))
            # 否则，就拼接当前目录和目标目录，得到新的目录
            else:
                # 如果目标目录不以/结尾，就加上/
                if not target_dir.endswith('/'):
                    target_dir += '/'
                current_dir = os.path.join(current_dir, target_dir)
            # 如果新的目录存在，就发送一个成功的响应给客户端
            if os.path.exists(current_dir):
                response = 'OK ' + current_dir
            # 否则，就发送一个失败的响应给客户端
            else:
                response = '目录不存在'
            # 发送响应给客户端
            client_sock.send(response.encode())
            # 返回新的当前目录
            return current_dir


    # 客服端读取服务端文件的方法
    def send_file(self, client_sock, command, current_dir):
        # 把命令分割为两部分，第一部分是get，第二部分是文件名
        parts = command.split(' ', 1)
        if len(parts) >= 2:
            _, filename = parts
        else:
            # 如果没有文件名，发送错误响应
            client_sock.send('ERROR: 没有指定文件名'.encode())
            return
        # 拼接当前目录和文件名，得到文件的完整路径
        filepath = os.path.join(current_dir, filename)
        # 如果文件存在，就发送一个成功的响应给客户端，包括文件名和文件大小
        if os.path.exists(filepath):
            try:
                filesize = os.path.getsize(filepath)
                response = f'OK {filesize} {filename}'
                client_sock.send(response.encode())
                # 打开文件，准备读取数据
                while True:
                    data = client_sock.recv(BUFFER_SIZE).decode()
                    if data == "start send file":
                        break
                with open(filepath, 'rb') as f:
                    # 初始化已发送的字节数为0
                    sent = 0
                    # 如果断点不为0，就从断点处开始读取数据
                    if self.breakpoint != 0:
                        # 移动文件指针到断点处
                        f.seek(self.breakpoint)
                        # 从断点处开始累加已发送的字节数
                        sent = self.breakpoint
                    # 循环读取数据，直到文件发送完毕
                    while sent < filesize:
                        # 读取数据
                        data = f.read(BUFFER_SIZE)
                        if not data:
                            break
                        # 发送数据
                        client_sock.send(data)
                        # 累加已发送的字节数
                        sent += len(data)
                    # 在控制台打印发送完成的消息
                    self.log_signal.emit('发送完成：' + filename)
            except Exception as e:
                self.log_signal.emit(f'发送异常: {str(e)}')
                response = '发送异常'
                client_sock.send(response.encode())
        # 否则，就发送一个失败的响应给客户端
        else:
            response = '文件不存在'
            client_sock.send(response.encode())

    # 接收文件并保存的方法
    def receive_file(self, client_sock, command, current_dir):
        self.lock.acquire()
        try:
            # 把命令分割为三部分，第一部分是put，第二部分是文件大小，第三部分是文件名
            _, filesize_str, filename = command.split(' ', 2)
            filesize = int(filesize_str)
            # 用os.path.basename函数来提取出文件名
            base_filename = os.path.basename(filename)
            # 拼接当前目录和文件名，得到文件的完整路径
            filepath = os.path.join(current_dir, base_filename)
            self.log_signal.emit(current_dir)
            self.log_signal.emit(base_filename)
            self.log_signal.emit(filepath)
            # 发送一个成功的响应给客户端，包括文件名
            response = 'OK ' + filename
            client_sock.send(response.encode())
            # 以追加模式或写入模式打开文件
            with open(filepath, 'ab' if self.breakpoint != 0 else 'wb') as f:
                # 从断点处开始累加已接收的字节数
                received = self.breakpoint
                self.size_flag = 0
                # 循环接收数据，直到文件接收完毕
                while received < filesize:
                    # 接收数据
                    data = client_sock.recv(BUFFER_SIZE)
                    if not data:
                        break
                    # 去除数据头大小内容
                    if self.size_flag == 0:
                        data = data[str(filesize).__len__():]
                        self.size_flag = 1
                    # 写入数据
                    f.write(data)
                    # 累加已接收的字节数
                    received += len(data)
                    if received >= filesize:
                        self.set_breakpoint(client_sock, 'restart 0')
                        break
            # 在控制台打印接收完成的消息
            self.log_signal.emit('接收完成：' + filename)
        except Exception as e:
            # 打印详细的异常信息
            self.log_signal.emit(f'接收异常: {str(e)}')
        finally:
            self.lock.release()

    # 设置断点的方法
    def set_breakpoint(self, client_sock, command):
        # 获取断点的位置
        parts = command.split(' ', 1)
        if len(parts) >= 2:
            _, self.breakpoint = parts
            self.breakpoint = int(self.breakpoint)
        else:
            # 如果没有断点位置，默认为0
            self.breakpoint = 0
        # 发送断点给客户端
        response = str(self.breakpoint)
        client_sock.send(response.encode())

    # 将用户添加到数据库中的方法
    def add_user_to_database(self, client_sock, command, db):
        # 从命令中分离出用户名和密码
        username, password, role = command.split(' ')[1:]
        # 调用DBManager对象的insert_user方法，插入用户到数据库
        result = db.insert_user(username, password, role)
        # 如果结果为True，表示插入成功，发送一个成功的响应给客户端
        if result:
            client_sock.send(f"OK 注册成功".encode())
        # 否则，表示插入失败，发送一个失败的响应给客户端
        else:
            client_sock.send(f"ERROR 用户名已存在 ${command}".encode())


class ServerThread(QThread):
    log_signal = pyqtSignal(str)

    def __init__(self, ip, port):
        super().__init__()
        self.ip = ip
        self.port = port
        self.server = None
        self.running = False

    def run(self):
        try:
            self.server = FTPServer(self.ip, self.port)
            self.server.log_signal.connect(self.handle_log)
            self.running = True
            self.log_signal.emit("服务器已启动")
            self.server.start()
        except Exception as e:
            self.log_signal.emit(f"启动失败: {str(e)}")

    def stop(self):
        if self.server:
            self.server.server_sock.close()
            self.running = False
            self.log_signal.emit("服务器已停止")

    def handle_log(self, message):
        self.log_signal.emit(message)


class FTPGUI(QtWidgets.QWidget):
    def __init__(self):
        super().__init__()
        self.server_thread = None
        self.init_ui()

    def init_ui(self):
        self.setWindowTitle('FTP 服务器')
        self.setFixedSize(800, 600)

        # IP 地址选择
        ip_layout = QtWidgets.QHBoxLayout()
        ip_label = QtWidgets.QLabel('服务器IP:')
        self.ip_combo = QtWidgets.QComboBox()
        self.ip_combo.addItems(self.get_local_ips())
        ip_layout.addWidget(ip_label)
        ip_layout.addWidget(self.ip_combo)

        # 端口设置
        port_layout = QtWidgets.QHBoxLayout()
        port_label = QtWidgets.QLabel('端口号:')
        self.port_edit = QtWidgets.QLineEdit('8888')
        port_layout.addWidget(port_label)
        port_layout.addWidget(self.port_edit)

        # 控制按钮
        self.start_btn = QtWidgets.QPushButton('启动服务')
        self.stop_btn = QtWidgets.QPushButton('停止服务')
        self.stop_btn.setEnabled(False)

        # 日志区域
        self.log_area = QtWidgets.QTextEdit()
        self.log_area.setReadOnly(True)

        # 布局管理
        main_layout = QtWidgets.QVBoxLayout()
        control_layout = QtWidgets.QHBoxLayout()

        control_layout.addLayout(ip_layout)
        control_layout.addLayout(port_layout)
        control_layout.addWidget(self.start_btn)
        control_layout.addWidget(self.stop_btn)

        main_layout.addLayout(control_layout)
        main_layout.addWidget(self.log_area)

        self.setLayout(main_layout)

        # 连接信号
        self.start_btn.clicked.connect(self.start_server)
        self.stop_btn.clicked.connect(self.stop_server)

    def get_local_ips(self):
        # 初始化 IP 列表
        ip_list = []
        # 获取所有网卡的地址信息
        net_if_addrs = psutil.net_if_addrs()
        for interface, addrs in net_if_addrs.items():
            for addr in addrs:
                # 筛选 IPv4 地址
                if addr.family == socket.AF_INET:
                    ip_list.append(addr.address)
        return ip_list

    def start_server(self):
        ip = self.ip_combo.currentText()
        try:
            port = int(self.port_edit.text())
        except ValueError:
            self.log_area.append("错误：端口号必须为数字！")
            return

        # 创建服务器线程
        self.server_thread = ServerThread(ip, port)
        self.server_thread.log_signal.connect(self.update_log)
        self.server_thread.start()

        # 更新按钮状态
        self.start_btn.setEnabled(False)
        self.stop_btn.setEnabled(True)
        self.ip_combo.setEnabled(False)
        self.port_edit.setEnabled(False)

    def stop_server(self):
        if self.server_thread:
            self.server_thread.stop()
            self.server_thread.quit()

            self.start_btn.setEnabled(True)
            self.stop_btn.setEnabled(False)
            self.ip_combo.setEnabled(True)
            self.port_edit.setEnabled(True)

    def update_log(self, message):
        self.log_area.append(f"[{QtCore.QDateTime.currentDateTime().toString('hh:mm:ss')}] {message}")

    def closeEvent(self, event):
        self.stop_server()
        event.accept()


if __name__ == '__main__':
    app = QtWidgets.QApplication(sys.argv)
    window = FTPGUI()
    window.show()
    sys.exit(app.exec_())